package com.yike.websocket.util;

import lombok.extern.log4j.Log4j2;

import javax.websocket.SendHandler;
import javax.websocket.SendResult;
import javax.websocket.Session;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author Sakura
 * @date 2024/9/13 11:40
 */
@Log4j2
public class WebsocketUtil {

    private static final Map<String, Session> ONLINE_SESSION = new ConcurrentHashMap<>();

    /**
     * 添加session
     */
    public static Session getSession(String userId){
        return ONLINE_SESSION.get(userId);
    }

    /**
     * 添加session
     */
    public static void addSession(String userId, Session session){
        // 一个用户只允许一个session链接
        //ONLINE_SESSION.putIfAbsent(userId, session);
        ONLINE_SESSION.put(userId, session);
        log.info("User [{}] connected. Total online users: {}", userId, ONLINE_SESSION.size());
    }

    /**
     * 移除session
     */
    public static void removeSession(String userId){
        ONLINE_SESSION.remove(userId);
        log.info("User [{}] disconnected. Total online users: {}", userId, ONLINE_SESSION.size());
    }

    /**
     * 给单个用户推送消息
     */
    public static void sendMessage(String userId, String message){
        Session session = ONLINE_SESSION.get(userId);
        if(session == null){
            log.warn("Session for user [{}] not found", userId);
            return;
        }
        sendMessage(session, message);
    }

    public static void sendMessage(Session session, String message) {
        if (session != null) {
            session.getAsyncRemote().sendText(message, new SendHandler() {
                @Override
                public void onResult(SendResult result) {
                    if (result.isOK()) {
                        log.info("消息发送成功");
                    } else {
                        log.warn("消息发送失败: " + result.getException());
                    }
                }
            });
            // session.getAsyncRemote().sendText(message);
        }
    }

    /**
     * 给所有用户发消息
     */
    public static void sendMessageForAll(String message) {
        ONLINE_SESSION.forEach((userId, session) -> {
            CompletableFuture.runAsync(() -> sendMessage(session, message))
                    .exceptionally(ex -> {
                        log.error("Failed to send message to user [{}]: {}", userId, ex.getMessage());
                        return null;
                    });
        });
    }

    /**
     * 给指定的多个用户推送消息
     */
    public static void sendMessageForUsers(Set<String> userIds, String message) {
        userIds.forEach(userId -> {
            Session session = ONLINE_SESSION.get(userId);
            if (session == null) {
                log.warn("Session for user [{}] not found", userId);
                return;
            }
            CompletableFuture.runAsync(() -> sendMessage(session, message))
                    .exceptionally(ex -> {
                        log.error("Failed to send message to user [{}]: {}", userId, ex.getMessage());
                        return null;
                    });
        });
    }

}